<#
.NOTES
    Author:     Lior Arviv, MSFT
    Version:    2.0.0
    Created:    01/02/2018
    Updated:    14/12/2020 - Added support for Azure Arc enabled servers

#>

#Requires -Modules Az.Accounts, Az.ResourceGraph

if (-not (Get-AzContext)) {
    Write-Host "Please authenticate to Azure using 'Connect-AzAccount'"
}

$query = @"
securityresources
| where type == 'microsoft.security/assessments' and name == 'ffff0522-1e88-47fc-8382-2a80ba848f5d'
| extend status = properties.status.code, resourceid = properties.resourceDetails.Id
| where status == 'Unhealthy'
| project resourceid
"@

$vms = Search-AzGraph -Query $query
foreach ($vm in $vms) {
    # Check if the resource is a regular virtual machine or Azure Arc connected
    if (($vm.resourceid -split '\/')[-3] -match "Microsoft.Compute") {
        $vmName = ($vm.resourceid -split '\/')[-1]
        Write-Host "Working on $vmName" -ForegroundColor Green
        $vmStatus = Invoke-AzRestMethod -Path ('{0}/instanceView?api-version=2020-06-01' -f $vm.resourceid) -Method GET |
        Select-Object -ExpandProperty Content | ConvertFrom-Json | Select-Object -ExpandProperty Statuses
        if ($vmStatus.displayStatus -match 'VM running') {
            $res = Invoke-AzRestMethod -Path ('{0}/providers/Microsoft.Security/serverVulnerabilityAssessments/default?api-Version=2015-06-01-preview' -f $vm.resourceid) -Method PUT
            if ($res.StatusCode -notmatch '200|202') {
                Write-Host ($res.Content | ConvertFrom-Json).Error.message -ForegroundColor Red
            }
        }
        else {
            Write-Host "$vmName is currently stopped. Skipping this one" -ForegroundColor Yellow
        }
    }
    else {
        $vmName = ($vm.resourceid -split '\/')[-1]
        Write-Host "Working on $vmName" -ForegroundColor Green
        $vmStatus = Invoke-AzRestMethod -Path ('{0}?api-version=2019-12-12' -f $vm.resourceid) -Method GET |
        Select-Object -ExpandProperty Content | ConvertFrom-Json | Select-Object -ExpandProperty properties
        if ($vmStatus.status -match 'Connected') {
            $res = Invoke-AzRestMethod -Path ('{0}/providers/Microsoft.Security/serverVulnerabilityAssessments/default?api-Version=2015-06-01-preview' -f $vm.resourceid) -Method PUT
            if ($res.StatusCode -notmatch '200|202') {
                Write-Host ($res.Content | ConvertFrom-Json).Error.message -ForegroundColor Red
            }
        }
        else {
            Write-Host "$vmName is currently stopped. Skipping this one" -ForegroundColor Yellow
        }
    } 
}